# -*- coding:utf-8 -*-
import math
import collections


class LanguageModel(object):
    def __init__(self, lm_dict_path, alpha=0.5):
        self._alpha = alpha
        self._word_coll_dict = self.load_word_coll_dict(lm_dict_path)
        pass

    @staticmethod
    def load_word_coll_dict(path):
        word_coll_dict = dict()
        with open(path, 'r', encoding='utf-8') as fr:
            for i, line in enumerate(fr):
                if i == 0:
                    continue
                ss = line.strip().split('\t')
                word_coll_dict[ss[0]] = float(ss[1])
        print('load {} items from {} over !'.format(len(word_coll_dict), path))
        return word_coll_dict

    def score(self, query: list, candidate: list) -> float:
        """
        calculate matching score by language model
        :param query: query sentence, list of words
        :param candidate: candidate sentence, list of words
        :return: score after normalization
        """
        if len(query) == 0 or len(candidate) == 0:
            return 0.0
        score1 = self._score(query, query)
        score2 = self._score(query, candidate)
        return (score2 - score1) / len(query)
        pass

    def _score(self, query, candidate) -> float:
        """
        match without normalization
        """
        factor = 1.0 / len(candidate)
        freq_dict = collections.defaultdict(int)

        for word in candidate:
            freq_dict[word] += factor

        score = 0
        for word in query:
            pwd = freq_dict[word] if word in freq_dict else 0.0
            pwc = self._word_coll_dict[word] if word in self._word_coll_dict else 1e-9

            tmp = (1.0 - self._alpha) * pwd + self._alpha * pwc
            # if tmp > 0:
            score += math.log(tmp)
        return score


class TranslationModel(LanguageModel):
    def __init__(self, lm_dict_path, trans_dict_path, alpha=0.2):
        super(TranslationModel, self).__init__(lm_dict_path, alpha)
        self._w2w_trans_dict = self.load_w2w_trans_dict(trans_dict_path)
        pass

    @staticmethod
    def load_w2w_trans_dict(trans_dict_path):
        w2w_trans_dict = collections.defaultdict(dict)
        count = 0
        with open(trans_dict_path, 'r', encoding='utf-8') as fr:
            for line in fr:
                ss = line.rstrip().split('\t')
                # if len(ss[0]) < 2 or len(ss[1]) < 2:
                #    continue
                count += 1
                word1, word2, prob = ss[0], ss[1], float(ss[2])
                w2w_trans_dict[word1][word2] = prob
        print('load {} word2word translation probabilities from {} over ...'.format(count, trans_dict_path))
        return w2w_trans_dict
        pass

    def get_word2word_trans_prob(self, word1, word2):
        if word1 not in self._w2w_trans_dict:
            return 0.0
        if word2 not in self._w2w_trans_dict[word1]:
            return 0.0
        return self._w2w_trans_dict[word1][word2]

    def _score(self, query: list, candidate: list):
        factor = 1.0 / len(candidate)
        freq_dict = collections.defaultdict(int)

        for word in candidate:
            freq_dict[word] += factor

        score = 0
        for word1 in query:
            ptr = 0.0
            for word2 in candidate:
                if word1 == word2:
                    ptr += 1.0 * factor
                else:
                    ptr += self.get_word2word_trans_prob(word1, word2) * factor
            pwc = self._word_coll_dict[word1] if word1 in self._word_coll_dict else 1e-9

            tmp = (1.0 - self._alpha) * ptr + self._alpha * pwc
            score += math.log(tmp)
        return score


if __name__ == '__main__':
    lm = LanguageModel(lm_dict_path='../lib/word.coll.dict')
    sent1 = '淘宝 不 支持 花呗 吗'.split()
    sent2 = '支付 方式 不 支持 花呗'.split()
    print(lm.score(sent1, sent2))
    print(lm.score(sent2, sent1))

    sent3 = '花呗 能 在 淘宝 上用 不'.split()
    print(lm.score(sent1, sent3))

    tr = TranslationModel(lm_dict_path='../lib/word.coll.dict', trans_dict_path="../lib/post2cmnt.trans.dict")
    sent1 = '淘宝 不 支持 花呗 吗'.split()
    sent2 = '支付 方式 不 支持 花呗'.split()
    print(tr.score(sent1, sent2))
    print(tr.score(sent2, sent1))

    sent3 = '花呗 能 在 淘宝 上用 不'.split()
    print(tr.score(sent1, sent3))

    pass
